#include "RSA.h"

#include <algorithm>

#include  <assert.h>

static int qpow(int base, int idx, int mod){
    int ret = 1;
    while(idx){
        if(idx&1){
            ret *= base; ret %= mod;
            idx ^= 1;
        }
        base *= base; base %= mod;
        idx >>= 1;
    }
    return ret;
}

static bool is_prime(int num){
    if(num < 2){
        return false;
    }
    if(num == 2){
        return true;
    }
    for(int i = 2; i * i <= num; ++i){
        if(num % i == 0){
            return false;
        }
    }
    return true;
}

RSA::RSA(int p, int q)
    :p_(p), q_(q), e_(-1), d_(-1){
    assert(is_prime(p));
    assert(is_prime(q));

    for(int i = std::max(p_, q_)+1;; ++i){
        if(is_prime(i)){
            e_ = i;
            break;
        }
    }

    for(int k = 1;; ++k){
        if((k*fhi_n()+1)%e() == 0){
            d_ = (k*fhi_n()+1)/e();
            break;
        }
    }

}

int RSA::n() const{
    return p_ * q_;
}

int RSA::fhi_n() const{
    return (p_-1)*(q_-1);
}

int RSA::e()const{
    return e_;
}

int RSA::d()const{
    return d_;
}

std::vector<int> RSA::encode(const std::string& message) const{
    std::vector<int> ret;
    for(char c : message){
        ret.push_back(qpow(c, e(), n()));
    }
    return ret;
}

std::string RSA::decode(const std::vector<int>& message) const{
    std::string ret;
    for(int num : message){
        ret.push_back(char(qpow(num, d(), n())));
    }
    return ret;
}